#include "local_planner/local_planner.h"
#include <pcl_conversions/pcl_conversions.h>
#include "transform_utils/type_transform.h"

namespace Planner
{
    LocalPlanner::LocalPlanner(const ros::NodeHandle &node_handle)
        : nh_(node_handle)
    {
        init();
    }

    void LocalPlanner::init()
    {
        int collision_checker_iter_num;
        double hangle, vangle;
        nh_.param("local_planner/collision_checker_iter_num", collision_checker_iter_num, 5);
        nh_.param("local_planner/hangle", hangle, 90.0);
        nh_.param("local_planner/vangel", vangle, 90.0);
        nh_.param("local_planner/secure_dis", secure_dis_, 1.0);
        nh_.param("local_planner/check_duration", check_duration_, 0.5);
        nh_.param("local_planner/type_", type_, 0);

        int bspline_order;
        double des_gap;
        nh_.param("local_planner/bspline_order", bspline_order, 3);
        nh_.param("local_planner/des_gap", des_gap, 1.5);

        collision_checker_.setParams(secure_dis_, collision_checker_iter_num, hangle, vangle);
        traj_opt_.setParams(nh_);
        bspline_smooth_.setParams(bspline_order, des_gap);

        cloud_map_ptr_ = boost::make_shared<pcl::PointCloud<pcl::PointXYZ>>(pcl::PointCloud<pcl::PointXYZ>());
        kdtree_ptr_ = boost::make_shared<pcl::KdTreeFLANN<pcl::PointXYZ>>(pcl::KdTreeFLANN<pcl::PointXYZ>());

        string cloud_topic, odom_topic;
        nh_.param("local_planner/cloud_topic", cloud_topic, string("velodyne_points"));
        nh_.param("local_planner/odom_topic", odom_topic, std::string("/m100/odometry_sensor1/odometry"));

        cloud_map_sub_ = nh_.subscribe(cloud_topic, 1, &LocalPlanner::cloudMapCallback, this);
        odom_sub_ = nh_.subscribe(odom_topic, 1, &LocalPlanner::odomCallback, this);

        replan_server_ = nh_.advertiseService("replan_check", &LocalPlanner::replanCheckCallback, this);
    }

    void LocalPlanner::cloudMapCallback(const sensor_msgs::PointCloud2 &msgs)
    {
        // 对于rotors仿真器需要进行tf变换
        sensor_msgs::PointCloud2 new_msg = msgs;
        if (type_ == 0)
        {
            listener_.waitForTransform("/world", "/m100/m100/velodyne", ros::Time(0), ros::Duration(0.05));
            listener_.lookupTransform("/world", "/m100/m100/velodyne", ros::Time(0), transform_);
            pcl_ros::transformPointCloud("/world", msgs, new_msg, listener_);
        }

        std::lock_guard<std::mutex> gurad(mutex_);
        cloud_map_ptr_.reset(new pcl::PointCloud<pcl::PointXYZ>);
        kdtree_ptr_.reset(new pcl::KdTreeFLANN<pcl::PointXYZ>);

        pcl::fromROSMsg(new_msg, *cloud_map_ptr_);
        kdtree_ptr_->setInputCloud(cloud_map_ptr_);
    }

    void LocalPlanner::odomCallback(const nav_msgs::Odometry &msgs)
    {
        std::lock_guard<std::mutex> gurad(mutex_);
        curr_pose_[0] = msgs.pose.pose.position.x;
        curr_pose_[1] = msgs.pose.pose.position.y;
        curr_pose_[2] = msgs.pose.pose.position.z;
        curr_orien_.x() = msgs.pose.pose.orientation.x;
        curr_orien_.y() = msgs.pose.pose.orientation.y;
        curr_orien_.z() = msgs.pose.pose.orientation.z;
        curr_orien_.w() = msgs.pose.pose.orientation.w;
    }

    bool LocalPlanner::replanCheckCallback(quad_msgs::ReplanCheck::Request &request,
                                           quad_msgs::ReplanCheck::Response &response)
    {
        ROS_INFO("Receive local_planner request");
        {
            std::lock_guard<std::mutex> gurad(mutex_);
            collision_checker_.setInputCloud(cloud_map_ptr_, kdtree_ptr_);
            collision_checker_.setCurrPose(curr_pose_, curr_orien_);
            traj_opt_.setCloudMap(cloud_map_ptr_, kdtree_ptr_);
        }

        if (!request.obtain_goal && request.polynomials.size() < 2)
        {
            ROS_ERROR("not enough polynomials in request");
            response.replan_flag = quad_msgs::ReplanCheck::Response::REPLAN_FAIL;
            return true;
        }

        // 依次判断当前滑动窗口内的分段轨迹安全性
        ROS_INFO("check sampled points in sliding window");
        vector<int> unsafe_traj_id;
        vector<Vector3d> waypoints;
        bool flag;

        // 检测过后，若输入的判断轨迹为i,i+1,i+2段，则输出的waypoints包含第i,i+1,i+2,i+3个点，unsafe_traj_id表示存在风险的轨迹编号
        if (request.obtain_goal)
            flag = sampleWaypointsCheck(request, 0, request.polynomials.size() - 1, unsafe_traj_id, waypoints);
        else
            flag = sampleWaypointsCheck(request, 0, request.polynomials.size() - 2, unsafe_traj_id, waypoints);

        if (!flag)
        {
            response.replan_flag = quad_msgs::ReplanCheck::Response::REPLAN_FAIL;
            return true;
        }

        if (unsafe_traj_id.empty())
        {
            response.replan_flag = quad_msgs::ReplanCheck::Response::SAFE;
            return true;
        }

        // 根据重规划的路径点生成局部轨迹
        ROS_INFO("replanning the collision part");

        // 从replan_pts中选取其中需要进行重规划的部分
        quad_msgs::State start_state = request.waypoints[unsafe_traj_id.front()];
        Vector3d start_pos(start_state.position.x, start_state.position.y, start_state.position.z);
        Vector3d start_vel(start_state.velocity.x, start_state.velocity.y, start_state.velocity.z);
        Vector3d start_acc(start_state.acceleration.x, start_state.acceleration.y, start_state.acceleration.z);

        // 除非request中包含终点且最后一段轨迹存在风险，否则goal_state的选取为需重规划轨迹的下一段
        bool revise_end = (request.obtain_goal && unsafe_traj_id.back() == request.polynomials.size() - 1);

        Vector3d goal_vel, goal_acc;
        vector<Vector3d> collision_pts;
        if (revise_end)
        {
            goal_vel.setZero();
            goal_acc.setZero();
            collision_pts.insert(collision_pts.end(), waypoints.begin() + unsafe_traj_id.front(), waypoints.end());
        }
        else
        {
            quad_msgs::State goal_state;
            goal_state = request.waypoints[unsafe_traj_id.back() + 2];
            goal_vel = Vector3d(goal_state.velocity.x, goal_state.velocity.y, goal_state.velocity.z);
            goal_acc = Vector3d(goal_state.acceleration.x, goal_state.acceleration.y, goal_state.acceleration.z);

            for (int i = unsafe_traj_id.front(); i <= unsafe_traj_id.back() + 1; i++)
                collision_pts.push_back(waypoints[i]);
            collision_pts.push_back(Vector3d(goal_state.position.x, goal_state.position.y, goal_state.position.z));
        }

        // 局部轨迹优化
        // 除非是包含终点且最后一段轨迹存在风险的情况，采用上述方法简化后collision_pts至少会有三个
        // TODO 这种情况暂时先做刹车处理，之后再考虑其它方式
        if (collision_pts.size() < 3)
        {
            cout << "collision_pts:\n";
            for (auto coll_pt: collision_pts)
                cout << coll_pt[0] << " " << coll_pt[1] << " " << coll_pt[2] << endl;
                
            response.replan_flag = quad_msgs::ReplanCheck::Response::REPLAN_FAIL;
            return true; 
        }
        else
        {
            ROS_INFO("local trajectory optimization");
            traj_opt_.setWayPoints(collision_pts);
            traj_opt_.setBoundaryCondition(start_vel, start_acc, goal_vel, goal_acc);
            int result = traj_opt_.optimize();

            if (result == -1) // -1为nlopt::FAILURE
            {
                response.replan_flag = quad_msgs::ReplanCheck::Response::REPLAN_FAIL;
                return true;
            }

            ROS_INFO("add replanning result to response");
            vector<MatrixXd> coeffs;
            vector<double> time_seq;
            traj_opt_.getCoeffs(coeffs);
            traj_opt_.getTimeSeq(time_seq);

            // 将优化结果写入response中
            if (revise_end)
            {
                // 记路径点为p0\p1\p2\p3，若p1-p2\p2-p3间存在碰撞风险，则front=1,back=2,轨迹优化的区间为p1-p3，
                // 则start_id=1，end_id=3(还需加上request中的id编号)，需更新的waypoints为p2\p3，polynomials为p1-p2\p2-p3
                response.start_id = unsafe_traj_id.front() + request.waypoint_id.front();
                response.end_id = unsafe_traj_id.back() + 1 + request.waypoint_id.front();
                coeffsToTrajMsg(coeffs, time_seq, response.polynomials);
                for (size_t i = 0; i < time_seq.size(); i++)
                {
                    quad_msgs::State state = getStateFromCoeffs(response.polynomials[i].coeffs_x, response.polynomials[i].coeffs_y,
                                                                response.polynomials[i].coeffs_z, response.polynomials[i].duration);
                    response.waypoints.push_back(state);
                }
                response.replan_flag = quad_msgs::ReplanCheck::Response::MOVE_END;
            }
            else
            {
                // 记路径点为p0\p1\p2\p3\p4，若p1-p2\p2-p3间存在碰撞风险，则front=1,back=2,轨迹优化的区间为p1-p4，
                // 则start_id=1，end_id=4(还需加上request中的id编号)，需更新的waypoints为p2\p3，polynomials为p1-p2\p2-p3\p3-p4
                response.start_id = unsafe_traj_id.front() + request.waypoint_id.front();
                response.end_id = unsafe_traj_id.back() + 2 + request.waypoint_id.front();
                coeffsToTrajMsg(coeffs, time_seq, response.polynomials);
                for (size_t i = 0; i < time_seq.size() - 1; i++)
                {
                    quad_msgs::State state = getStateFromCoeffs(response.polynomials[i].coeffs_x, response.polynomials[i].coeffs_y,
                                                                response.polynomials[i].coeffs_z, response.polynomials[i].duration);
                    response.waypoints.push_back(state);
                }
                response.replan_flag = quad_msgs::ReplanCheck::Response::REPLAN_SUCCESS;
            }
        }
        return true;
    }

    bool LocalPlanner::sampleWaypointsCheck(quad_msgs::ReplanCheck::Request &request, int start, int end,
                                            vector<int> &unsafe_traj_id, vector<Vector3d> &waypoints)
    {
        unsafe_traj_id.clear();
        waypoints.clear();
        vector<double> coeffsx, coeffsy, coeffsz;
        Vector3d last_pt;
        double obs_dis;

        last_pt[0] = request.waypoints[start].position.x;
        last_pt[1] = request.waypoints[start].position.y;
        last_pt[2] = request.waypoints[start].position.z;
        if (!collision_checker_.isFeasible(last_pt, obs_dis))
        {
            ROS_ERROR("Collision occured at replan start position (%f, %f, %f), obstacle gap is %f",
                      last_pt[0], last_pt[1], last_pt[2], obs_dis);
            return false;
        }
        waypoints.push_back(last_pt);

        // 1. i表示路径点编号还是轨迹段数编号，traj_server和local_planner中是否混淆?
        //    i在traj_server中表示路径点编号,下方for循环中i表示轨迹编号，unsafe_traj_id存储的也为轨迹id
        // 2. 当sample_pts.size() < 3时该如何处理?
        //    目前的代码结构存在一些问题，例如需要判断最后一个路径点的安全性，
        //    这给整个规划过程带来了许多的边界条件，因此做一些简化处理，最后一个路径点默认为安全的，
        //    这样能够保证始终有一个无需变动的轨迹优化终止状态
        for (int i = start; i <= end; i++)
        {
            coeffsx = request.polynomials[i].coeffs_x;
            coeffsy = request.polynomials[i].coeffs_y;
            coeffsz = request.polynomials[i].coeffs_z;
            double duration = request.polynomials[i].duration;
            // 对分段轨迹进行离散采样及碰撞检测
            int segs = std::ceil(duration / check_duration_);
            cout << i << ": " << duration << "s, " << segs << " segments\n";
            for (int j = 1; j <= segs; j++)
            {
                // 计算当前轨迹点的位置及方向
                Vector3d pt = getPosFromPieceCoeffs(coeffsx, coeffsy, coeffsz, duration / segs * j);
                Vector3d direct;
                direct = (pt - last_pt).normalized();

                Vector3d replan_pt;
                bool safety = collision_checker_.isFeasible(pt, obs_dis);
                ROS_INFO("Sample point at segment %d, position: (%f, %f, %f), obstacle gap: %f",
                         i, pt[0], pt[1], pt[2], obs_dis);

                if (safety)
                {
                    if (j == segs)
                        waypoints.push_back(pt);
                    continue;
                }

                // 轨迹点存在碰撞风险，需要进行重规划
                if (unsafe_traj_id.empty() || unsafe_traj_id.back() != i)
                    unsafe_traj_id.push_back(i);

                if (!collision_checker_.replanByTree(pt, direct, replan_pt))
                {
                    ROS_ERROR("Replan failed at trajectory point (%f, %f, %f)", pt[0], pt[1], pt[2]);
                    return false;
                }
                last_pt = pt;
                if (j == segs)
                    waypoints.push_back(replan_pt);
            }
        }
        return true;
    }
} // namespace Planner